#!/usr/bin/env python
# coding=utf-8
from DBUtils import GetCveList, GetCpeByCve, GetSourcePackageList, GetCpeList
import lxml
from lxml import etree

source_packages = GetSourcePackageList()

def compareName(pkg, repo_name):
        p1 = pkg.replace(".","-").split('-')
        p2 = repo_name.replace(".","-").split('-')
        if len(p1)>1:
            for k in range(len(p1)):
                if p1[k].startswith("lib"):
                    p1[k] = p1[k][3:]
                if p1[k] in ["haskell", "java", "python", "perl", "el", "js"]:
                    p1[k] = ""
        if len(p2)>1:
            for k in range(len(p2)):
                if p2[k].startswith("lib"):
                    p2[k] = p2[k][3:]
                if p2[k] in ["haskell", "java", "python", "perl", "el", "js"]:
                    p2[k] = ""
    
        p1 = ''.join(p1).lower()
        p2 = ''.join(p2).lower()
        # print p1
        # print p2
        return p1 == p2

def getSoftwareNameFromCpe(cpe):
    parts = cpe.split(":")
    assert parts[0]=="cpe"
    if parts[1]!="/a":
        return None
    vendor, software = None, None
    if len(parts) == 5:
        vendor, software = parts[2], parts[3]
    else:
        software = parts[3]
    return software

def GetPackageBySoftware(software):
    candidates = []
    for sp, in source_packages:
        if compareName(sp, software):
            candidates.append(sp)
    if not candidates:
        return None
    return software, candidates

    
def CountCpeNumber():
    '''
    统计所有CVE中包含CPE信息的个数
    '''
    cve_count = 0
    cpe_count = 0
    for cve, in GetCveList():
        cve_count += 1
        CPE_list = GetCpeByCve(cve)
        if CPE_list:
            cpe_count += 1

    print cve_count, cpe_count

def InitCpePackageMap():
    '''
    建立CPE与Package之间的映射
    '''
    # @Depracated
    # cve_count = 0
    # for cve, in GetCveList():
    #     # if cve_count>=10:
    #         # break
    #     result = []
    #     CPE_list = GetCpeByCve(cve)
    #     softwares = [getSoftwareNameFromCpe(cpe) for cpe, in CPE_list]
    #     softwares = [x for x in softwares if x is not None]
    #     softwares = list(set(softwares))
    #     for sw in softwares:
    #         ret = GetPackageBySoftware(sw)
    #         if ret:
    #             result.append(ret)
    #     if result:
    #         print cve, 
    #         for item in result:
    #             print "%s:%s" % (item[0], ','.join(item[1]))
    #         cve_count+=1


    cpe_list = [cpe for cpe, in GetCpeList()]
    for pkg, in GetSourcePackageList():
        match_software_list = []
        softwares = [getSoftwareNameFromCpe(cpe) for cpe in cpe_list]
        softwares = [x for x in softwares if x]
        softwares = list(set(softwares))
        for sw in softwares:
            if compareName(pkg, sw):
                match_software_list.append(sw)
        print pkg, match_software_list
# InitCpePackageMap()
# print compareName("s-el", "ssh")

def GetSoftwareByCve(cve):
    '''
    从DB中获取CVE关联的CPE列表，再从CPE列表中获得一个软件列表，去重后返回list
    '''
    vul_sw = [getSoftwareNameFromCpe(x) for x, in GetCpeByCve(cve)]
    vul_sw = [x for x in vul_sw if x is not None]
    vul_sw = list(set(vul_sw))
    return vul_sw

